import { ref } from 'vue';

export class ComfyUIService {
    private ws: WebSocket | null = null;
    private queue: Array<{ resolve: Function; reject: Function }> = [];
    public connected = ref(false);
    public progress = ref(0);
    public status = ref('');

    constructor(private wsUrl: string) { }

    connect() {
        if (this.ws?.readyState === WebSocket.OPEN) return;

        this.ws = new WebSocket(this.wsUrl);

        this.ws.onopen = () => {
            this.connected.value = true;
            this.status.value = '已连接';
        };

        this.ws.onclose = () => {
            this.connected.value = false;
            this.status.value = '连接已断开';
            setTimeout(() => this.connect(), 5000);
        };

        this.ws.onerror = (error) => {
            console.error('WebSocket error:', error);
            this.status.value = '连接错误';
        };

        this.ws.onmessage = (event) => {
            const data = JSON.parse(event.data);
            this.handleMessage(data);
        };
    }

    private handleMessage(data: any) {
        switch (data.type) {
            case 'progress':
                this.progress.value = data.data.value;
                this.status.value = `生成中...${Math.round(data.data.value * 100)}%`;
                break;
            case 'executed':
                this.status.value = '生成完成';
                if (this.queue.length > 0) {
                    const { resolve } = this.queue.shift()!;
                    resolve(data.data);
                }
                break;
            case 'error':
                this.status.value = '生成失败';
                if (this.queue.length > 0) {
                    const { reject } = this.queue.shift()!;
                    reject(new Error(data.data.message));
                }
                break;
        }
    }

    async generateImage(workflow: any, params: any): Promise<string> {
        if (!this.connected.value) {
            throw new Error('WebSocket 未连接');
        }

        return new Promise((resolve, reject) => {
            this.queue.push({ resolve, reject });
            this.ws?.send(JSON.stringify({
                type: 'execute',
                data: {
                    workflow,
                    params
                }
            }));
        });
    }

    disconnect() {
        this.ws?.close();
    }
}

// 创建单例
export const comfyUIService = new ComfyUIService(
    import.meta.env.VITE_COMFYUI_WS_URL
); 